#ifndef __ctpl_stl_thread_pool_H__
#define __ctpl_stl_thread_pool_H__

#include <functional>
#include <thread>
#include <atomic>
#include <vector>
#include <memory>
#include <exception>
#include <future>
#include <mutex>
#include <queue>

namespace cptl {
    namespace detail {
        template <typename T>
        class Queue {
        public:
            bool push(T const & value) {
                std::unique_lock<std::mutex> lock(mutex_);
                q_.push(value);
                return true;
            }

            // deletes the retrieved element, do not use for non integral types
            bool pop(T & v) {
                std::unique_lock<std::mutex> lock(mutex_);
                if (q_.empty()) {
                    return false;
                }

                v = q_.front();
                q_.pop();
                return true;
            }

            bool empty() {
                std::unique_lock<std::mutex> lock(mutex_);
                return q_.empty();
            }
        private:
            std::queue<T> q_;
            std::mutex mutex_;
        };
    }

    class ThreadPool {
    public:
        ThreadPool() {
            Init();
        }

        ThreadPool(int nThreads) {
            this->Init();
            this->Resize(nThreads);
        }

        ~ThreadPool() {
            Stop(true);
        }

        // wait for all computing threads to finish and stop all threads
        // may be called asynchronously to not pause the calling thread while waiting
        // if isWait == true, all the functions in the queue are run, otherwise the queue is 
        // cleared without running the functions
        void Stop(bool isWait = false) {
            if (!isWait) {
                if (this->isStop) {
                    return;
                }
                this->isStop = true;
                for (int i = 0, n = this->Size(); i < n; ++i) {
                    //command the threads to stop
                    *this->flags[i] = true;
                }
                //empty the queue
                this->ClearQueue();
            }
            else {
                if (this->isDone || this->isStop) {
                    return;
                }
                // give the waiting threads a command to finish
                this->isDone = true;
            }

            {
                std::unique_lock<std::mutex> lock(this->mutex);
                // stop all waiting threads
                this->cv.notify_all();
            }

            // wait for the computing threads to finish
            for (int i = 0; i < static_cast<int>(this->threads.size()); ++i) {
                if (this->threads[i]->joinable()) {
                    this->threads[i]->join();
                }
            }

            // if there were no threads in the pool but some functors in the queue, the
            // functors are not deleted by the threads. therefore delete them here
            this->ClearQueue();
            this->threads.clear();
            this->flags.clear();
        }

        // get the number of running threads in the pool
        int Size() { return static_cast<int>(this->threads.size()); }

        // number of idle threads
        int n_idle() { return this->nWaiting; }

        std::thread & GetThread(int i) { return *this->threads[i]; }

        // change the number of threads in the pool
        // should be called from one thread, otherwise be careful to not interleave,
        // also with this->stop().
        // nTreads must be >= 0
        void Resize(int nThreads) {
            if (!this->isStop && !this->isDone) {
                int oldThreads = static_cast<int>(this->threads.size());
                // if the number of threads is increased
                if (oldThreads <= nThreads) {
                    this->threads.resize(nThreads);
                    this->flags.resize(nThreads);

                    for (int i = oldThreads; i < nThreads; ++i) {
                        this->flags[i] = std::make_shared<std::atomic<bool>>(false);
                        this->SetThread(i);
                    }
                }
                // the number of threads is decreased
                else {
                    for (int i = oldThreads - 1; i >= nThreads; --i) {
                        //this thread will finish
                        *this->flags[i] = true;
                        this->threads[i]->detach();
                    }

                    {
                        // stop the detached threads that were waiting
                        std::unique_lock<std::mutex> lock(this->mutex);
                        this->cv.notify_all();
                    }

                    // safe to delete because the threads are detached
                    this->threads.resize(nThreads);

                    // safe to delete because the threads hava copies of 
                    // shared_ptr of the flags, not originals
                    this->flags.resize(nThreads);
                }
            }
        }

        // empty the queue
        void ClearQueue() {
            std::function<void(int id)> * _f;
            while (this->q.pop(_f)) {
                delete _f;
            }
        }

        // pop a functional wrapper to the original function
        std::function<void(int)> Pop() {
            std::function<void(int)> * _f = nullptr;
            this->q.pop(_f);

            // at return, delete the function even if an exception occurred
            std::unique_ptr<std::function<void(int id)>> func(_f);
            std::function<void(int)> f;
            if (_f)
                f = *_f;
            return f;
        }

        template <typename F, typename... Rest>
        auto push(F && f, Rest && ... rest) -> std::future<decltype(f(0, rest...))> {
            auto pck = std::make_shared<std::packaged_task<decltype(f(0, rest...))(int)>>(
                std::bind(std::forward<F>(f), std::placeholders::_1, std::forward<Rest>(rest)...)
            );

            auto _f = new std::function<void(int id)>([pck](int id) {
                (*pck)(id);
            });

            this->q.push(_f);
            std::unique_lock<std::mutex> lock(this->mutex);
            this->cv.notify_one();
            return pck->get_future();
        }

        // run the user's function that excepts argument int - id of the running thread.
        // returned value is templatized operator returns std::future, where the user 
        // can get the result and rethrow the catched exceptions
        template<typename F>
        auto push(F && f) -> std::future<decltype(f(0))> {
            auto pck = std::make_shared<std::packaged_task<decltype(f(0))(int)>>(std::forward<F>(f));
            auto _f = new std::function<void(int id)> ([pck](int id) {
                (*pck)(id);
            });
            this->q.push(_f);
            std::unique_lock<std::mutex> lock(this->mutex);
            this->cv.notify_one();
            return pck->get_future();
        }

    private:
        // deleted
        ThreadPool(const ThreadPool &) = delete;
        ThreadPool(ThreadPool &&) = delete;
        ThreadPool & operator = (const ThreadPool &) = delete;
        ThreadPool & operator = (ThreadPool &&) = delete;

        void Init() {
            nWaiting = 0;
            isStop = false;
            isDone = false;
        }

        void SetThread(int i) {
            std::shared_ptr<std::atomic<bool>> flag(this->flags[i]);
            auto f = [this, i, flag /* a copy of the shared ptr to the flag */] () { 
                // use reference to keep trace the thread's flag in the flags array
                std::atomic<bool> & _flag = *flag; 
                std::function<void(int id)> * _f;
                bool isPop = this->q.pop(_f);

                while (true) {
                    while (isPop) { // if there is anything in the queue

                        //at return, delete the function even if an exception occurred
                        std::unique_ptr<std::function<void(int id)>> func(_f);
                        (*_f)(i);

                        // the thread is wanted to stop, return even if the queue is not empty yet
                        if (_flag) {
                            return;
                        }
                        else {
                            isPop = this->q.pop(_f);
                        }
                    }

                    // the queue is empty here, wait for the next command
                    std::unique_lock<std::mutex> lock(this->mutex);
                    ++this->nWaiting;
                    this->cv.wait(lock, [this, &_f, &isPop, &_flag]() {
                        isPop = this->q.pop(_f);
                        return isPop || this->isDone || _flag;
                    });
                    --this->nWaiting;

                    // if the queue is empty and this->isDone == true or *flag then return
                    if (!isPop) {
                        return;
                    }
                }
            };
            this->threads[i].reset(new std::thread(f));// Can I replace new with std::make_unique<> ?
        }

    private:
        std::vector<std::unique_ptr<std::thread>> threads;
        std::vector<std::shared_ptr<std::atomic<bool>>> flags;
        detail::Queue<std::function<void(int id)> *> q;
        std::atomic<bool> isDone;
        std::atomic<bool> isStop;
        std::atomic<int> nWaiting; // how many threads are waiting

        std::mutex mutex;
        std::condition_variable cv;
    };
}

#endif